﻿// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using Microsoft.AspNetCore.Http;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Threading.RateLimiting;
using System.Timers;

namespace FlyingEye.Middlewares.RateLimiting
{
    public class UserPartitionedRateLimiter
    {
        public PartitionedRateLimiter<HttpContext> PartitionedRateLimiter { get; }

        public ConcurrentUpDownCounter Counter { get; }

        public UserPartitionedRateLimiter(PartitionedRateLimiter<HttpContext> partitionedRateLimiter,
            ConcurrentUpDownCounter counter)
        {
            PartitionedRateLimiter = partitionedRateLimiter;
            Counter = counter;
        }
    }

    internal class RateLimitingMetrics : IDisposable
    {
        private Func<PartitionedRateLimiter<HttpContext>> _factory;

        private Timer _timer;

        /// <summary>
        /// 用户 id 映射到分区限速器
        /// </summary>
        private readonly ConcurrentDictionary<Guid, UserPartitionedRateLimiter> _userLimiterMaps = new();

        public RateLimitingMetrics(Func<PartitionedRateLimiter<HttpContext>> factory, TimeSpan interval)
        {
            _factory = factory;
            _timer = new Timer(interval.TotalMilliseconds);
            _timer.Elapsed += TimerElapsed;
            _timer.Start();
        }

        private void TimerElapsed(object? sender, ElapsedEventArgs e)
        {
            // 如果 http 请求个数为 0，则移除限速器
            // 计数为 0 的限速器，一定是调用了 LeaseEnd 方法后才会为 0
            // 所以在这里移除计数为 0 的限制器是线程安全的
            _userLimiterMaps.RemoveAll(item => item.Value.Counter.Count <= 0);
        }

        /// <summary>
        /// 租聘开始
        /// </summary>
        public PartitionedRateLimiter<HttpContext> LeaseStart(Guid userId)
        {
            if (_userLimiterMaps.TryGetValue(userId, out var userLimiter))
            {
                userLimiter.Counter.Increment();
                return userLimiter.PartitionedRateLimiter;
            }
            else
            {
                userLimiter = new UserPartitionedRateLimiter(_factory.Invoke(), new ConcurrentUpDownCounter());
                userLimiter.Counter.Increment();

                // 这个 userId key 一定是不存在的，所以不必验证返回结果
                _userLimiterMaps.TryAdd(userId, userLimiter);

                return userLimiter.PartitionedRateLimiter;
            }
        }

        /// <summary>
        /// 租聘结束
        /// </summary>
        public void LeaseEnd(Guid userId)
        {
            if (_userLimiterMaps.TryGetValue(userId, out var limiter))
            {
                limiter.Counter.Decrement();
            }
        }

        public void Dispose()
        {
            _userLimiterMaps.Clear();
            _timer.Stop();
            _timer.Close();
        }
    }
}
